from ChatBot import Bot
from ai_chat import chat
import re

bot = None
# 记录机器人开关
open_dict = {}
# 开启机器人关键词
open_key_word = '召唤小夏AI'
# 关闭小夏AI
close_key_word = '关闭小夏AI'
# 切换用户角色身份
change_role_key_word = '切换角色'

# 记录用户角色身份,key 用户昵称,value 身份描述
user_role_dict = {}


def handle_message(name_data, message_data):
    for msg in message_data:
        if msg[0] == 'SYS':
            continue
        name = msg[0]
        message = msg[1]
        print(f'{name}发来了消息:{message}')
        # 判断是否开启关键词
        if message == open_key_word:
            open_dict[name] = 'open'
            bot.send_message(name, f'小夏已开启AI助手模式')
            return
        # 判断是否命中关闭关键词
        if message == close_key_word:
            open_dict[name] = 'close'
            bot.send_message(name, f'小夏已关闭AI助手模式')
            return
        # 判断是否命中切换角色关键词
        if message.startswith(change_role_key_word):
            role_desc = message[len(change_role_key_word) + 1:]
            user_role_dict[name] = role_desc
            bot.send_message(name, f'已切换角色:{role_desc}')
            return

        # 先判断机器人开关是否开启
        if open_dict.get(name) is None:
            bot.send_message(name,
                             f'亲爱的{name},小夏正在忙率中...,是否要开启小夏AI?是请回:{open_key_word}，否请回:{close_key_word},切换角色身份请回:{change_role_key_word}，角色描述，例如回复 切换角色，请以一个高冷帅哥的口吻和我聊天')
            return

        if open_dict[name] == 'open':
            res = chat(message, user_role_dict.get(name))
            bot.send_message(name, res)


nicknames = ['岁月如歌', '程嘉慧', '洁怡的小月牙']
bot = Bot(nicknames, handle_message)
bot.hold()
